package com.xp.ai.embedding;


import com.google.protobuf.Api;
import com.xp.ai.util.ApiKey;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;

import java.util.List;

/***
 * 客服问答系统相似度匹配系统 demo
 * @author xp
 */
public class QuestionVectorDemo {

    public static void main(String[] args) {
        //向量数据库
        EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
                .host("localhost")
                .port(5433)
                .user("root")
                .password("123456")
                .database("postgres")
                .table("question_embedding")
                .dimension(1024)
                .useIndex(true)
                .indexListSize(100)
                .createTable(true)
                .dropTableFirst(false)
                .build();

        //创建向量化的模型
        OpenAiEmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl(ApiKey.GJ_BASE_URL)
                .apiKey(ApiKey.GJ_API_KEY)
                .modelName(ApiKey.GJ_EMBEDDING_MODEL)
                .build();


        //只有第一次的时候需要添加，第二次之后就不需要添加向量数据了
       /* //字符串片段
        TextSegment jrText = TextSegment.from("今日头条1000粉丝才可以面基");
        TextSegment hsText = TextSegment.from("火山视频100粉丝就可以面基了");
        TextSegment dyText = TextSegment.from("抖音1000粉丝就可以使用电脑直播了");
        TextSegment wdText = TextSegment.from("我的名字叫Tz");
        Response<Embedding> mjResponse = embeddingModel.embed(jrText);
        Response<Embedding> hsResponse = embeddingModel.embed(hsText);
        Response<Embedding> dyResponse = embeddingModel.embed(dyText);
        Response<Embedding> tzResponse = embeddingModel.embed(wdText);

        //加入向量数据库
        embeddingStore.add(mjResponse.content(), jrText);
        embeddingStore.add(hsResponse.content(), hsText);
        embeddingStore.add(dyResponse.content(), dyText);
        embeddingStore.add(tzResponse.content(), wdText);*/

        //提问
        Response<Embedding> questionRes = embeddingModel.embed("火山视频的面基要求是多少粉丝？");
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(questionRes.content())
                .minScore(0.0)
                .maxResults(10)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        searchResult.matches().forEach(match -> {
            System.out.println("检索文本为： " + match.embedded().text());
            System.out.println("文本相似度： " + match.score());
        });


    }
}
